# 环节5：检索 - 根据用户问题找到最相关的文档片段
# 功能：接收用户问题，生成向量后在Milvus中检索最相关的文本片段
# 依赖：需先完成环节1-4构建向量数据库

import time
from pymilvus import Collection, connections, utility
from langchain_community.embeddings import DashScopeEmbeddings

# 配置参数（与环节1-4保持一致）
CONFIG = {
    "tongyi_api_key": "sk-6434b554122c4b049ceb805d703f695b",  # 你的阿里云API Key
    "milvus_host": "localhost",
    "milvus_port": "19530",
    "collection_name": "insurance_manual",
    "top_k": 3  # 检索返回的最相关片段数量
}


def initialize_embeddings(api_key):
    """初始化通义千问嵌入模型"""
    try:
        embeddings = DashScopeEmbeddings(
            model="text-embedding-v4",
            dashscope_api_key=api_key
        )
        # 验证模型是否可用
        embeddings.embed_query("测试")
        return embeddings
    except Exception as e:
        raise Exception(f"嵌入模型初始化失败: {str(e)}")


def connect_milvus(host, port):
    """连接到Milvus服务"""
    if not connections.has_connection("default"):
        connections.connect(
            alias="default",
            host=host,
            port=port
        )

    # 检查集合是否存在
    if not utility.has_collection(CONFIG["collection_name"]):
        raise Exception(f"集合 {CONFIG['collection_name']} 不存在，请先运行环节1-4构建数据库")

    # 加载集合
    collection = Collection(CONFIG["collection_name"])
    collection.load()
    return collection


def get_vector_dimension(collection):
    """获取向量维度（用于验证查询向量）"""
    vector_field = next(f for f in collection.schema.fields if f.name == "vector")
    return vector_field.params["dim"]


def retrieve_relevant_docs(collection, query, embeddings, top_k):
    """
    根据用户问题检索最相关的文档片段
    参数：
        collection - Milvus集合对象
        query - 用户问题文本
        embeddings - 嵌入模型
        top_k - 返回的最相关片段数量
    返回：
        包含相关片段信息的列表，每个元素包含页码、内容和相似度
    """
    # 1. 生成查询向量
    query_vector = embeddings.embed_query(query)

    # 2. 验证向量维度
    vector_dim = get_vector_dimension(collection)
    if len(query_vector) != vector_dim:
        raise ValueError(f"查询向量维度错误: 应为{vector_dim}维，实际{len(query_vector)}维")

    # 3. 配置检索参数
    search_params = {
        "metric_type": "L2",  # 与存储时使用的度量方式一致
        "params": {"nprobe": 10}  # 检索精度参数，值越大精度越高但速度越慢
    }

    # 4. 执行检索
    start_time = time.time()
    results = collection.search(
        data=[query_vector],  # 传入查询向量列表
        anns_field="vector",  # 检索的向量字段名
        param=search_params,
        limit=top_k,  # 返回的结果数量
        output_fields=["text", "page"]  # 需要返回的其他字段
    )
    search_time = time.time() - start_time

    # 5. 处理检索结果
    relevant_docs = []
    for hit in results[0]:  # results是列表，第一个元素是当前查询的结果
        relevant_docs.append({
            "id": hit.id,
            "page": hit.entity.page,  # 文档页码
            "text": hit.entity.text,  # 文本内容
            "distance": hit.distance,  # 与查询的距离（值越小越相关）
            "similarity": 1 / (1 + hit.distance)  # 转换为相似度（0-1之间）
        })

    print(f"🔍 检索完成（耗时{search_time:.2f}秒）")
    print(f"   问题: {query}")
    print(f"   返回最相关的{len(relevant_docs)}个片段\n")

    return relevant_docs


def display_results(relevant_docs):
    """格式化显示检索结果"""
    for i, doc in enumerate(relevant_docs, 1):
        print(f"【相关片段 {i}】")
        print(f"页码: 第{doc['page']}页")
        print(f"相似度: {doc['similarity']:.2%}")
        print(f"内容: {doc['text'][:200]}...")  # 只显示前200字符
        print("-" * 80)


def main():
    """检索环节主函数"""
    try:
        # 1. 初始化组件
        print("===== 初始化检索组件 =====")
        embeddings = initialize_embeddings(CONFIG["tongyi_api_key"])
        collection = connect_milvus(CONFIG["milvus_host"], CONFIG["milvus_port"])
        print(f"✅ 组件初始化完成，向量维度: {get_vector_dimension(collection)}")

        # 2. 接收用户问题（实际应用中可改为输入交互）
        user_queries = [
            "平安福2024版的投保条件是什么？",
            "健康告知需要提供哪些信息？",
            "保险理赔的流程是怎样的？"
        ]

        # 3. 对每个问题执行检索并显示结果
        for query in user_queries:
            relevant_docs = retrieve_relevant_docs(
                collection=collection,
                query=query,
                embeddings=embeddings,
                top_k=CONFIG["top_k"]
            )
            display_results(relevant_docs)

        # 4. 释放资源
        collection.release()
        print("\n===== 检索流程结束 =====")

    except Exception as e:
        print(f"\n❌ 检索失败: {str(e)}")


if __name__ == "__main__":
    main()
